/** 
 * Copyright 2016 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 *
 * Licensed under the Amazon Software License (the "License"). You may not use this file 
 * except in compliance with the License. A copy of the License is located at
 *
 *   http://aws.amazon.com/asl/
 *
 * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, 
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, express or implied. See the License for the 
 * specific language governing permissions and limitations under the License.
 */
package com.amazon.alexa.avs.auth.companionapp;

import com.amazon.alexa.avs.auth.AccessTokenListener;
import com.amazon.alexa.avs.config.DeviceConfig;
import com.amazon.alexa.avs.config.DeviceConfig.CompanionAppInformation;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.sql.Date;
import java.util.Map;
import java.util.Timer;
import java.util.TimerTask;
import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;

/**
 * Entry points for dealing with authentication and provisioning. Owns exchanging credentials for
 * tokens and managing those tokens.
 */
public class CompanionAppAuthManager {

    private static final Logger log = LoggerFactory.getLogger(CompanionAppAuthManager.class);

    /**
     * How many times to retry exchanging refreshToken for an accessToken.
     */
    private static final int TOKEN_REFRESH_RETRY_COUNT = 3;

    /**
     * How long in seconds before trying again to exchange refreshToken for an accessToken.
     */
    private static final int TOKEN_REFRESH_RETRY_INTERVAL_IN_S = 2;

    /**
     * A map from sessionId to codeVerifier.
     *
     * New sessionId and codeVerifiers are generated by {@link #getDeviceProvisioningInfo()} and
     * mapped together for easy recall in the future.
     */
    private final Map<String, String> sessionIdToCodeVerifier =
            new ConcurrentHashMap<String, String>();

    /**
     * Device configuration information.
     */
    private final DeviceConfig deviceConfig;

    /**
     * Client for exchanging credentials (authCode, clientId, refreshToken, etc) for accessTokens.
     */
    private final OAuth2ClientForPkce pkceOAuth2Client;

    /**
     * Handles generating codeVerifier, codeChallenge, and the codeChallengeMethod.
     */
    private final CodeChallengeWorkflow codeChallengeWorkflow;

    /**
     * The current tokens being used to make requests to AVS.
     */
    private OAuth2TokensForPkce tokens;

    private final AccessTokenListener accessTokenListener;

    private final Timer refreshTimer;

    /**
     * Creates an {@link CompanionAppAuthManager} object.
     *
     * @param deviceConfig
     * @param oAuth2Client
     * @param codeChallengeWorkflow
     * @param accessTokenListener
     */
    public CompanionAppAuthManager(DeviceConfig deviceConfig, OAuth2ClientForPkce oAuth2Client,
            CodeChallengeWorkflow codeChallengeWorkflow, AccessTokenListener accessTokenListener) {
        this.deviceConfig = deviceConfig;
        this.pkceOAuth2Client = oAuth2Client;
        this.codeChallengeWorkflow = codeChallengeWorkflow;
        this.accessTokenListener = accessTokenListener;
        this.refreshTimer = new Timer();

        if (deviceConfig.getCompanionAppInfo() != null
                && deviceConfig.getCompanionAppInfo().getClientId() != null
                && deviceConfig.getCompanionAppInfo().getRefreshToken() != null) {
            this.refreshTimer.schedule(new RefreshTokenTimerTask(), 0);
        }
    }

    /**
     * Return a {@link DeviceProvisioningInfo} populated with the necessary codeChallenge
     * information and device information, including productId and dsn.
     *
     * @return The information necessary to start the device provisioning process.
     */
    public DeviceProvisioningInfo getDeviceProvisioningInfo() {
        codeChallengeWorkflow.generateProofKeyParameters();

        // Get everything that we need from the CodeChallengeWorkflow.
        String codeChallenge = codeChallengeWorkflow.getCodeChallenge();
        String codeChallengeMethod = codeChallengeWorkflow.getCodeChallengeMethod();
        String codeVerifier = codeChallengeWorkflow.getCodeVerifier();
        String sessionId = UUID.randomUUID().toString();

        // Map sessionId back to codeVerifier so that we can retrieve it later given a sessionId.
        sessionIdToCodeVerifier.put(sessionId, codeVerifier);

        // Return the object will all the necessary information that can be serialized by the client
        // or server later.
        DeviceProvisioningInfo deviceProvisioningInfo =
                new DeviceProvisioningInfo(deviceConfig.getProductId(), deviceConfig.getDsn(),
                        sessionId, codeChallenge, codeChallengeMethod);
        return deviceProvisioningInfo;
    }

    /**
     * Requests accessToken and refreshToken from LWA by exchanging information provided by the
     * companion, and the codeVerifier, for tokens.
     *
     * @param companionProvisioningInfo
     *            The information provided by the companion application or service.
     * @throws IOException
     *             If an I/O exception occurs.
     */
    public void exchangeCompanionInfoForTokens(
            CompanionAppProvisioningInfo companionProvisioningInfo) throws IOException {
        // Pull out all information from the companion app
        String sessionId = companionProvisioningInfo.getSessionId();
        String clientId = companionProvisioningInfo.getClientId();
        String authCode = companionProvisioningInfo.getAuthCode();
        String redirectUri = companionProvisioningInfo.getRedirectUri();
        String codeVerifier = sessionIdToCodeVerifier.get(sessionId);

        // If we're unable to pull a valid codeVerifier from the map of sessionId->codeVerifier,
        // then the passed sessionId is invalid
        if (codeVerifier == null) {
            throw new InvalidSessionIdException(sessionId);
        }

        // Exchange the authCode and codeVerifier for refreshToken and accessToken
        OAuth2TokensForPkce tokens = pkceOAuth2Client.exchangeAuthCodeForTokens(authCode,
                redirectUri, clientId, codeVerifier);
        setTokens(tokens);
    }

    /**
     * Set tokens returned from the {@link OAuth2ClientForPkce} where they need to go.
     *
     * @param tokens
     *            Retrieved from the {@link OAuth2ClientForPkce}.
     */
    private synchronized void setTokens(OAuth2TokensForPkce tokens) {
        this.tokens = tokens;

        CompanionAppInformation info = deviceConfig.getCompanionAppInfo();
        info.setClientId(tokens.getClientId());
        info.setRefreshToken(tokens.getRefreshToken());
        deviceConfig.saveConfig();

        refreshTimer.schedule(new RefreshTokenTimerTask(), new Date(tokens.getExpiresTime()));

        accessTokenListener.onAccessTokenReceived(tokens.getAccessToken());
    }

    /**
     * Exchanges a refreshToken for an accessToken.
     *
     * @throws IOException
     *             If an I/O exception occurs.
     */
    public void refreshTokens() throws IOException {
        if (deviceConfig.getCompanionAppInfo() != null) {
            String refreshToken = deviceConfig.getCompanionAppInfo().getRefreshToken();
            String clientId = deviceConfig.getCompanionAppInfo().getClientId();
            refreshTokens(refreshToken, clientId);
        }
    }

    /**
     * Exchanges a refreshToken for an accessToken.
     *
     * @param refreshToken
     *            The refreshToken.
     * @param clientId
     *            The clientId of the companion application/service used.
     * @throws IOException
     *             If an I/O exception occurs.
     */
    private void refreshTokens(String refreshToken, String clientId) throws IOException {
        OAuth2TokensForPkce tokens =
                pkceOAuth2Client.exchangeRefreshTokenForTokens(refreshToken, clientId);
        setTokens(tokens);
    }

    /**
     * @return the most recent tokens, or null.
     */
    public OAuth2TokensForPkce getTokens() {
        return tokens;
    }

    /**
     * @return whether or not there are tokens
     */
    public boolean hasTokens() {
        return tokens != null;
    }

    @SuppressWarnings("javadoc")
    public static class InvalidSessionIdException extends IllegalArgumentException {
        private static final long serialVersionUID = 1L;

        /**
         * @param sessionId
         *            the invalid sessionId
         */
        public InvalidSessionIdException(String sessionId) {
            super("The sessionId that you passed is incorrect or invalid: " + sessionId);
        }
    }

    /**
     * TimerTask for refreshing accessTokens every hour.
     */
    private class RefreshTokenTimerTask extends TimerTask {
        @Override
        public void run() {
            int tries = 0;
            while (tries < TOKEN_REFRESH_RETRY_COUNT) {
                try {
                    refreshTokens();
                    break;
                } catch (IOException e) {
                    try {
                        log.error(
                                "There was a problem connecting to the LWA service. Trying again in {} seconds",
                                TOKEN_REFRESH_RETRY_INTERVAL_IN_S);
                        Thread.sleep(TOKEN_REFRESH_RETRY_INTERVAL_IN_S);
                    } catch (InterruptedException ie) {
                        log.error("Interrupted while waiting to retry connecting to LWA", ie);
                    }
                    tries++;
                }
            }
        }
    }

}
